#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
符合生物学规律的脑波测试数据生成器

本模块生成用于测试过滤器和分析算法的真实脑波数据，
所有数据都基于人体生理学原理设计，确保状态转换的合理性。

作者: 脑波分析系统开发团队
创建时间: 2025年1月
版本: 1.0.0
"""

import numpy as np
import json
from datetime import datetime, timedelta
from typing import Dict, List, Tuple, Any
import math


class BiologicallyRealisticDataGenerator:
    """
    符合生物学规律的脑波数据生成器
    
    生成各种符合人体生理学的脑波状态转换序列，
    包括正常的睡眠-觉醒周期、运动后恢复等场景。
    """
    
    def __init__(self, sampling_rate: int = 512):
        """
        初始化数据生成器
        
        Args:
            sampling_rate: 采样率，默认512Hz
        """
        self.sampling_rate = sampling_rate
        self.frequency_bands = {
            'delta': (0.5, 4),    # Delta波：深睡眠
            'theta': (4, 8),      # Theta波：浅睡眠、冥想
            'alpha': (8, 13),     # Alpha波：放松、闭眼
            'beta': (13, 30),     # Beta波：清醒、专注
            'gamma': (30, 100)    # Gamma波：高度兴奋
        }
        
        # 定义各状态的典型频域特征（相对功率）
        self.state_profiles = {
            'deep_sleep': {'delta': 0.6, 'theta': 0.25, 'alpha': 0.1, 'beta': 0.04, 'gamma': 0.01},
            'light_sleep': {'delta': 0.35, 'theta': 0.4, 'alpha': 0.15, 'beta': 0.08, 'gamma': 0.02},
            'relaxed': {'delta': 0.15, 'theta': 0.25, 'alpha': 0.45, 'beta': 0.12, 'gamma': 0.03},
            'awake': {'delta': 0.1, 'theta': 0.15, 'alpha': 0.25, 'beta': 0.45, 'gamma': 0.05},
            'excited': {'delta': 0.05, 'theta': 0.1, 'alpha': 0.15, 'beta': 0.5, 'gamma': 0.2},
            'focused': {'delta': 0.08, 'theta': 0.12, 'alpha': 0.2, 'beta': 0.55, 'gamma': 0.05}
        }
        
        # 定义生物学上合理的状态转换路径
        self.valid_transitions = {
            'awake': ['relaxed', 'excited', 'focused'],
            'excited': ['awake', 'relaxed'],
            'focused': ['awake', 'relaxed'],
            'relaxed': ['awake', 'light_sleep', 'focused'],
            'light_sleep': ['relaxed', 'deep_sleep', 'awake'],
            'deep_sleep': ['light_sleep']
        }
        
        # 状态转换的典型时间（秒）
        self.transition_durations = {
            ('awake', 'relaxed'): (30, 120),
            ('relaxed', 'light_sleep'): (60, 300),
            ('light_sleep', 'deep_sleep'): (120, 600),
            ('deep_sleep', 'light_sleep'): (60, 180),
            ('light_sleep', 'relaxed'): (30, 120),
            ('relaxed', 'awake'): (15, 60),
            ('excited', 'relaxed'): (120, 300),
            ('awake', 'excited'): (10, 30),
            ('focused', 'relaxed'): (60, 180)
        }
    
    def generate_signal_for_state(self, state: str, duration: float, 
                                add_noise: bool = True, 
                                noise_level: float = 0.1) -> np.ndarray:
        """
        为指定状态生成脑波信号
        
        Args:
            state: 脑状态名称
            duration: 持续时间（秒）
            add_noise: 是否添加噪声
            noise_level: 噪声水平
            
        Returns:
            生成的信号数组
        """
        if state not in self.state_profiles:
            raise ValueError(f"未知状态: {state}")
        
        samples = int(duration * self.sampling_rate)
        t = np.linspace(0, duration, samples)
        signal = np.zeros(samples)
        
        profile = self.state_profiles[state]
        
        # 为每个频段生成信号分量
        for band, (low_freq, high_freq) in self.frequency_bands.items():
            if profile[band] > 0:
                # 在频段范围内随机选择几个频率
                num_components = np.random.randint(2, 5)
                for _ in range(num_components):
                    freq = np.random.uniform(low_freq, high_freq)
                    amplitude = profile[band] * np.random.uniform(0.5, 1.5)
                    phase = np.random.uniform(0, 2 * np.pi)
                    signal += amplitude * np.sin(2 * np.pi * freq * t + phase)
        
        # 添加生理性变化（缓慢的幅度调制）
        modulation_freq = np.random.uniform(0.01, 0.1)  # 0.01-0.1 Hz
        modulation = 1 + 0.3 * np.sin(2 * np.pi * modulation_freq * t)
        signal *= modulation
        
        # 添加噪声
        if add_noise:
            noise = np.random.normal(0, noise_level * np.std(signal), samples)
            signal += noise
        
        return signal
    
    def generate_state_transition(self, from_state: str, to_state: str, 
                                transition_duration: float = None) -> np.ndarray:
        """
        生成状态转换期间的信号
        
        Args:
            from_state: 起始状态
            to_state: 目标状态
            transition_duration: 转换持续时间（秒），如果为None则使用默认值
            
        Returns:
            转换期间的信号数组
        """
        # 验证转换的生物学合理性
        if to_state not in self.valid_transitions.get(from_state, []):
            raise ValueError(f"生物学上不合理的状态转换: {from_state} -> {to_state}")
        
        # 确定转换时间
        if transition_duration is None:
            transition_key = (from_state, to_state)
            if transition_key in self.transition_durations:
                min_dur, max_dur = self.transition_durations[transition_key]
                transition_duration = np.random.uniform(min_dur, max_dur)
            else:
                transition_duration = np.random.uniform(60, 180)  # 默认1-3分钟
        
        samples = int(transition_duration * self.sampling_rate)
        t = np.linspace(0, transition_duration, samples)
        
        # 生成渐变信号
        from_profile = self.state_profiles[from_state]
        to_profile = self.state_profiles[to_state]
        
        signal = np.zeros(samples)
        
        for band, (low_freq, high_freq) in self.frequency_bands.items():
            # 计算每个时间点的频段权重（线性插值）
            from_weight = from_profile[band]
            to_weight = to_profile[band]
            
            for i, time_point in enumerate(t):
                # 使用S型曲线进行平滑转换
                progress = i / (samples - 1)
                smooth_progress = 1 / (1 + np.exp(-10 * (progress - 0.5)))
                
                current_weight = from_weight + (to_weight - from_weight) * smooth_progress
                
                if current_weight > 0:
                    freq = np.random.uniform(low_freq, high_freq)
                    amplitude = current_weight * np.random.uniform(0.8, 1.2)
                    signal[i] += amplitude * np.sin(2 * np.pi * freq * time_point)
        
        # 添加转换期间的不稳定性
        instability = 0.2 * np.random.normal(0, 1, samples)
        signal += instability
        
        return signal
    
    def generate_long_term_stable_state(self, state: str, duration_minutes: int) -> Dict[str, Any]:
        """
        生成长期稳定状态的测试数据
        
        Args:
            state: 状态名称
            duration_minutes: 持续时间（分钟）
            
        Returns:
            包含信号数据和元信息的字典
        """
        duration_seconds = duration_minutes * 60
        signal = self.generate_signal_for_state(state, duration_seconds)
        
        return {
            'signal': signal.tolist(),
            'state': state,
            'duration_minutes': duration_minutes,
            'duration_seconds': duration_seconds,
            'sampling_rate': self.sampling_rate,
            'description': f'长期稳定{state}状态，持续{duration_minutes}分钟',
            'expected_classification': state,
            'test_type': 'long_term_stable'
        }
    
    def generate_normal_sleep_cycle(self) -> Dict[str, Any]:
        """
        生成正常的入睡过程：清醒 → 放松 → 浅睡 → 深睡
        
        Returns:
            包含完整睡眠周期数据的字典
        """
        states = ['awake', 'relaxed', 'light_sleep', 'deep_sleep']
        durations = [300, 600, 900, 1800]  # 5分钟、10分钟、15分钟、30分钟
        
        full_signal = []
        state_timeline = []
        current_time = 0
        
        for i, (state, duration) in enumerate(zip(states, durations)):
            # 生成稳定状态信号
            stable_signal = self.generate_signal_for_state(state, duration)
            full_signal.extend(stable_signal)
            
            state_timeline.append({
                'state': state,
                'start_time': current_time,
                'end_time': current_time + duration,
                'duration': duration
            })
            current_time += duration
            
            # 生成状态转换信号（除了最后一个状态）
            if i < len(states) - 1:
                next_state = states[i + 1]
                transition_signal = self.generate_state_transition(state, next_state)
                full_signal.extend(transition_signal)
                
                transition_duration = len(transition_signal) / self.sampling_rate
                state_timeline.append({
                    'state': f'transition_{state}_to_{next_state}',
                    'start_time': current_time,
                    'end_time': current_time + transition_duration,
                    'duration': transition_duration
                })
                current_time += transition_duration
        
        return {
            'signal': full_signal,
            'state_timeline': state_timeline,
            'total_duration': current_time,
            'sampling_rate': self.sampling_rate,
            'description': '正常入睡过程：清醒→放松→浅睡→深睡',
            'test_type': 'normal_sleep_cycle'
        }
    
    def generate_wake_up_cycle(self) -> Dict[str, Any]:
        """
        生成正常的醒来过程：深睡 → 浅睡 → 放松 → 清醒
        
        Returns:
            包含完整醒来周期数据的字典
        """
        states = ['deep_sleep', 'light_sleep', 'relaxed', 'awake']
        durations = [1200, 600, 300, 300]  # 20分钟、10分钟、5分钟、5分钟
        
        full_signal = []
        state_timeline = []
        current_time = 0
        
        for i, (state, duration) in enumerate(zip(states, durations)):
            stable_signal = self.generate_signal_for_state(state, duration)
            full_signal.extend(stable_signal)
            
            state_timeline.append({
                'state': state,
                'start_time': current_time,
                'end_time': current_time + duration,
                'duration': duration
            })
            current_time += duration
            
            if i < len(states) - 1:
                next_state = states[i + 1]
                transition_signal = self.generate_state_transition(state, next_state)
                full_signal.extend(transition_signal)
                
                transition_duration = len(transition_signal) / self.sampling_rate
                state_timeline.append({
                    'state': f'transition_{state}_to_{next_state}',
                    'start_time': current_time,
                    'end_time': current_time + transition_duration,
                    'duration': transition_duration
                })
                current_time += transition_duration
        
        return {
            'signal': full_signal,
            'state_timeline': state_timeline,
            'total_duration': current_time,
            'sampling_rate': self.sampling_rate,
            'description': '正常醒来过程：深睡→浅睡→放松→清醒',
            'test_type': 'wake_up_cycle'
        }
    
    def generate_exercise_recovery_cycle(self) -> Dict[str, Any]:
        """
        生成运动后恢复过程：清醒 → 兴奋 → 放松
        
        Returns:
            包含运动恢复周期数据的字典
        """
        states = ['awake', 'excited', 'relaxed']
        durations = [300, 600, 900]  # 5分钟、10分钟、15分钟
        
        full_signal = []
        state_timeline = []
        current_time = 0
        
        for i, (state, duration) in enumerate(zip(states, durations)):
            stable_signal = self.generate_signal_for_state(state, duration)
            full_signal.extend(stable_signal)
            
            state_timeline.append({
                'state': state,
                'start_time': current_time,
                'end_time': current_time + duration,
                'duration': duration
            })
            current_time += duration
            
            if i < len(states) - 1:
                next_state = states[i + 1]
                transition_signal = self.generate_state_transition(state, next_state)
                full_signal.extend(transition_signal)
                
                transition_duration = len(transition_signal) / self.sampling_rate
                state_timeline.append({
                    'state': f'transition_{state}_to_{next_state}',
                    'start_time': current_time,
                    'end_time': current_time + transition_duration,
                    'duration': transition_duration
                })
                current_time += transition_duration
        
        return {
            'signal': full_signal,
            'state_timeline': state_timeline,
            'total_duration': current_time,
            'sampling_rate': self.sampling_rate,
            'description': '运动后恢复过程：清醒→兴奋→放松',
            'test_type': 'exercise_recovery_cycle'
        }
    
    def generate_edge_case_data(self) -> List[Dict[str, Any]]:
        """
        生成边缘状态测试数据
        
        Returns:
            边缘状态测试数据列表
        """
        edge_cases = []
        
        # 放松与浅睡之间的边界状态
        relaxed_profile = self.state_profiles['relaxed']
        light_sleep_profile = self.state_profiles['light_sleep']
        
        # 创建中间状态
        boundary_profile = {}
        for band in relaxed_profile:
            boundary_profile[band] = (relaxed_profile[band] + light_sleep_profile[band]) / 2
        
        # 临时添加到状态配置中
        self.state_profiles['relaxed_light_sleep_boundary'] = boundary_profile
        
        signal = self.generate_signal_for_state('relaxed_light_sleep_boundary', 300)
        edge_cases.append({
            'signal': signal.tolist(),
            'description': '放松与浅睡之间的边界状态',
            'expected_classification': ['relaxed', 'light_sleep'],  # 可能的分类结果
            'test_type': 'edge_case',
            'edge_type': 'relaxed_light_sleep_boundary'
        })
        
        # 清理临时状态
        del self.state_profiles['relaxed_light_sleep_boundary']
        
        # 清醒与兴奋之间的边界状态
        awake_profile = self.state_profiles['awake']
        excited_profile = self.state_profiles['excited']
        
        boundary_profile = {}
        for band in awake_profile:
            boundary_profile[band] = (awake_profile[band] + excited_profile[band]) / 2
        
        self.state_profiles['awake_excited_boundary'] = boundary_profile
        signal = self.generate_signal_for_state('awake_excited_boundary', 300)
        edge_cases.append({
            'signal': signal.tolist(),
            'description': '清醒与兴奋之间的边界状态',
            'expected_classification': ['awake', 'excited'],
            'test_type': 'edge_case',
            'edge_type': 'awake_excited_boundary'
        })
        
        del self.state_profiles['awake_excited_boundary']
        
        return edge_cases
    
    def generate_extreme_state_data(self) -> List[Dict[str, Any]]:
        """
        生成极端状态测试数据
        
        Returns:
            极端状态测试数据列表
        """
        extreme_cases = []
        
        # 极度兴奋状态（增强beta和gamma波）
        extreme_excited_profile = self.state_profiles['excited'].copy()
        extreme_excited_profile['beta'] *= 1.5
        extreme_excited_profile['gamma'] *= 2.0
        
        self.state_profiles['extreme_excited'] = extreme_excited_profile
        signal = self.generate_signal_for_state('extreme_excited', 600)
        extreme_cases.append({
            'signal': signal.tolist(),
            'description': '极度兴奋状态（高强度运动后）',
            'expected_classification': 'excited',
            'test_type': 'extreme_case',
            'extreme_type': 'extreme_excited'
        })
        
        # 极度深睡状态（增强delta波）
        extreme_deep_sleep_profile = self.state_profiles['deep_sleep'].copy()
        extreme_deep_sleep_profile['delta'] *= 1.3
        extreme_deep_sleep_profile['theta'] *= 0.7
        
        self.state_profiles['extreme_deep_sleep'] = extreme_deep_sleep_profile
        signal = self.generate_signal_for_state('extreme_deep_sleep', 1800)
        extreme_cases.append({
            'signal': signal.tolist(),
            'description': '极度深睡状态（极度疲劳）',
            'expected_classification': 'deep_sleep',
            'test_type': 'extreme_case',
            'extreme_type': 'extreme_deep_sleep'
        })
        
        # 清理临时状态
        del self.state_profiles['extreme_excited']
        del self.state_profiles['extreme_deep_sleep']
        
        return extreme_cases
    
    def generate_comprehensive_test_dataset(self) -> Dict[str, Any]:
        """
        生成完整的测试数据集
        
        Returns:
            包含所有测试数据的综合数据集
        """
        dataset = {
            'metadata': {
                'generation_time': datetime.now().isoformat(),
                'sampling_rate': self.sampling_rate,
                'generator_version': '1.0.0',
                'description': '符合生物学规律的综合脑波测试数据集'
            },
            'long_term_stable_states': [],
            'state_transition_cycles': [],
            'edge_cases': [],
            'extreme_cases': []
        }
        
        # 生成长期稳定状态数据
        stable_states = [
            ('deep_sleep', 30),
            ('focused', 20),
            ('relaxed', 15)
        ]
        
        for state, duration in stable_states:
            data = self.generate_long_term_stable_state(state, duration)
            dataset['long_term_stable_states'].append(data)
        
        # 生成状态转换周期数据
        dataset['state_transition_cycles'].append(self.generate_normal_sleep_cycle())
        dataset['state_transition_cycles'].append(self.generate_wake_up_cycle())
        dataset['state_transition_cycles'].append(self.generate_exercise_recovery_cycle())
        
        # 生成边缘状态数据
        dataset['edge_cases'] = self.generate_edge_case_data()
        
        # 生成极端状态数据
        dataset['extreme_cases'] = self.generate_extreme_state_data()
        
        return dataset
    
    def save_dataset_to_file(self, dataset: Dict[str, Any], filename: str = None) -> str:
        """
        将数据集保存到文件
        
        Args:
            dataset: 要保存的数据集
            filename: 文件名，如果为None则自动生成
            
        Returns:
            保存的文件路径
        """
        if filename is None:
            timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
            filename = f'biologically_realistic_test_data_{timestamp}.json'
        
        filepath = f'd:\\python\\brain_wave\\test\\comprehensive_filter_analysis_test\\{filename}'
        
        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump(dataset, f, indent=2, ensure_ascii=False)
        
        return filepath


if __name__ == '__main__':
    # 生成测试数据集
    generator = BiologicallyRealisticDataGenerator()
    
    print("正在生成符合生物学规律的测试数据...")
    dataset = generator.generate_comprehensive_test_dataset()
    
    # 保存数据集
    filepath = generator.save_dataset_to_file(dataset)
    print(f"测试数据已保存到: {filepath}")
    
    # 打印数据集统计信息
    print("\n数据集统计信息:")
    print(f"长期稳定状态测试: {len(dataset['long_term_stable_states'])} 个")
    print(f"状态转换周期测试: {len(dataset['state_transition_cycles'])} 个")
    print(f"边缘状态测试: {len(dataset['edge_cases'])} 个")
    print(f"极端状态测试: {len(dataset['extreme_cases'])} 个")
    
    total_duration = 0
    for stable_state in dataset['long_term_stable_states']:
        total_duration += stable_state['duration_seconds']
    
    for cycle in dataset['state_transition_cycles']:
        total_duration += cycle['total_duration']
    
    print(f"总测试时长: {total_duration/60:.1f} 分钟")
    print("\n数据生成完成！")